package flinkdemo.faceSink;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.hadoop.fs.*;
import org.apache.hadoop.hdfs.DistributedFileSystem;

import java.net.URI;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

public class ImageDirSource extends RichParallelSourceFunction<ImageData> {
    private String imagePath = System.getenv("IMAGE_INPUT_PATH");
    private transient FileSystem fileSystem;
    private transient List<LocatedFileStatus> currentFileList;
    private transient LogTimestampPlugin plugin = null;
    private String pluginClassName;
    private String logPath;

    public ImageDirSource() {
        if (null == this.imagePath) {
            throw new RuntimeException("Not set image_input_path environmental variable");
        } else {
            System.out.println(String.format("Image input path %s", this.imagePath));
            this.pluginClassName = System.getenv("LOG_PLUGIN_CLASSNAME");
            if (null == this.pluginClassName) {
                this.pluginClassName = PrintLogPlugin.class.getName();
            }

            this.logPath = System.getenv("SCORE_LOG_PATH");
            if (null == this.logPath) {
                this.logPath = "";
            }

            System.out.println(String.format("Plugin %s %s", this.pluginClassName, this.logPath));
        }
    }

    @Override
    public void run(SourceContext<ImageData> sourceContext) throws Exception {
        System.out.println("wait 5s for load model.");
        Thread.sleep(5000L);
        Iterator var2 = this.currentFileList.iterator();

        while(var2.hasNext()) {
            LocatedFileStatus fileStatus = (LocatedFileStatus)var2.next();
            ImageData imageLabel = new ImageData();
            imageLabel.setId(fileStatus.getPath().getName());
            long fileLength = fileStatus.getLen();
            FSDataInputStream in = this.fileSystem.open(fileStatus.getPath());
            byte[] buffer = new byte[(int)fileLength];
            in.readFully(buffer);
            in.close();
            imageLabel.setImage(buffer);
            sourceContext.collect(imageLabel);
            if (null != this.plugin) {
                this.plugin.logTimestamp(imageLabel.getId(), System.currentTimeMillis(), "");
            }
        }

    }

    @Override
    public void cancel() {
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        System.out.println(String.format("Plugin %s %s", this.pluginClassName, this.logPath));
        int index = this.getRuntimeContext().getIndexOfThisSubtask();
        int taskNum = this.getRuntimeContext().getNumberOfParallelSubtasks();
        if (this.pluginClassName.isEmpty()) {
            this.plugin = null;
        } else {
            this.plugin = (LogTimestampPlugin)Class.forName(this.pluginClassName).newInstance();
        }

        if (null != this.plugin) {
            this.plugin.open(this.logPath, "source", index);
        }

        this.currentFileList = new ArrayList();
        Path imageRoot = new Path(this.imagePath);
        org.apache.hadoop.conf.Configuration hadoopConfig = new org.apache.hadoop.conf.Configuration();
        hadoopConfig.set("fs.hdfs.impl", DistributedFileSystem.class.getName());
        hadoopConfig.set("fs.file.impl", LocalFileSystem.class.getName());
        this.fileSystem = FileSystem.get(new URI(this.imagePath), hadoopConfig);
        RemoteIterator<LocatedFileStatus> it = this.fileSystem.listFiles(imageRoot, true);
        int i = 0;

        while(it.hasNext()) {
            LocatedFileStatus locatedFileStatus = (LocatedFileStatus)it.next();
            if (!locatedFileStatus.getPath().getName().startsWith(".")) {
                if (index == i % taskNum) {
                    this.currentFileList.add(locatedFileStatus);
                }
                ++i;
            }
        }

    }

    @Override
    public void close() throws Exception {
        if (null != this.plugin) {
            this.plugin.close();
        }

        if (null != this.fileSystem) {
            this.fileSystem.close();
        }

    }
}